import os
import sys
import md5
import time
import base64
import SocketServer
import shutil
import subprocess
import traceback
import locale
import string
from SimpleXMLRPCServer import SimpleXMLRPCServer

from config import DUMP_DIR, TIMEOUT, PM_PATH, PROCDUMP_PATH
from tricks import DETECTION_TRICKS, DEBUGGING_TRICKS
import report_filter

def doLog(msg):
	print "[%s] %s" % (time.ctime(), msg)

def timed_run(cmd, time_out = TIMEOUT, delta = 1):
	t0 = time.time()
	doLog("Executing %s" % (cmd))
	proc = subprocess.Popen(cmd.split(), stdout = subprocess.PIPE, stderr = subprocess.STDOUT)
	
	while time.time() < time_out + t0:
		time.sleep(delta)
		pstatus = proc.poll()
		if pstatus is not None:
			return pstatus
	
	proc.terminate()
	doLog("Kill %s" % (cmd))
	pstatus = proc.poll()
	return pstatus

def fileToBuf(filename):
	try:
		buf = file(filename, "rb").read()
		buf = base64.b64encode(buf)
		return buf
	except:
		return None

def checkDebuggingTricks(filename):
	ret = []

	with open(filename) as f:
		for line in f:
			for trick in DEBUGGING_TRICKS:
				info = line.split(',')
				for i in range(5):
					if info[i].lower().find(trick.lower()) > -1:
						if not trick in ret:
							doLog("Detected debugging trick %s (%s)" % (trick, DEBUGGING_TRICKS[trick]))
							ret.append(trick)
							ret.append(line)

	if not ret:
		doLog("No known debugging detection tricks found")
	
	return ret

def convert_char(char):
	if char in string.printable:
		return char
	else:
		return r'\x%02x' % ord(char)
		
def convert_to_printable(s):
	return ''.join([convert_char(c) for c in s])
	
def string2hex(st):
	return "".join("%02x" % ord(c) for c in st)
	
def checkDetectionTricks(filename):
	tricks = 0
	ret = []
	
	with open(filename, "rb") as f:
		buf = f.read()
		for trick in DETECTION_TRICKS:
			if buf.find(DETECTION_TRICKS[trick][::-1]) > -1:
				doLog("Detected VM trick %s (%s)" % (trick, convert_to_printable(DETECTION_TRICKS[trick])))
				ret.append(trick + " (%s)" % convert_to_printable(DETECTION_TRICKS[trick]))
				
	if not ret:
		doLog("No known VM detection tricks found")
		
	return ret
	
# Threaded mix-in
class AsyncXMLRPCServer(SocketServer.ThreadingMixIn,SimpleXMLRPCServer):
	pass

class CXmlRpcDumper:
	
	server = None
	_md5 = None
	busy = False

	def dump(self, filebuf, timeout=15):
		
		if self.busy:
			raise Exception("Server is busy")
		
		self.busy = True
		d = {}
		
		try:			
			if not os.path.exists(DUMP_DIR):
				os.mkdir(DUMP_DIR)
			filebuf = base64.b64decode(filebuf)
			self._md5 = md5.md5(filebuf).hexdigest()
			newDir = DUMP_DIR + os.sep + self._md5
			
			if os.path.exists(newDir):
				shutil.rmtree(newDir)
			os.mkdir(newDir)
			
			malwareFile = newDir + os.sep + "malware.exe"
			f = file(malwareFile, "wb")
			f.write(filebuf)
			f.close()
			
			pm_log = newDir + os.sep + "report.pml"
			pm_rpt = newDir + os.sep + "report.csv"
			os.system("start " + PM_PATH + " /Quiet /Minimized /Backingfile " + pm_log)
			os.system(PM_PATH + " /Waitforidle")
			
			# cmd = PROCDUMP_PATH + " -c 0 -s 3 -ma -w malware.exe " + newDir
			# proc = subprocess.Popen(cmd.split(), stdout = subprocess.PIPE, stderr = subprocess.STDOUT)
			
			try:
				d["report_warning"] = None
				timed_run(malwareFile, timeout)
			except WindowsError as err:
				msg = "WindowsError: %d, Message: " % (err.args[0])\
				+ err.args[1].decode(locale.getdefaultlocale()[1])
				doLog(msg)
				d["report_warning"] = msg
				
			doLog("Terminating ProcessMonitor")
			os.system(PM_PATH + " /Terminate")
			doLog("Saving csv report")
			os.system(PM_PATH + " /Openlog %s /SaveAs %s"% (pm_log, pm_rpt))
			d["report"] = fileToBuf(pm_rpt)
			
			doLog("Filtering report")
			filtered = report_filter.log_filter(pm_rpt)
			d["report_filtered"] = fileToBuf(filtered)
			doLog("Extracting signature")
			d["report_signature"] = fileToBuf(report_filter.extract_signature(filtered))
			d["tricks_debug"] = checkDebuggingTricks(filtered)
			
			# if proc.poll() is None:
				# proc.terminate()
			# file_list = os.listdir(newDir)
			# dump_file = newDir + os.sep + "malware.dmp"
			# for each in file_list:
				# if each.endswith(".dmp"):
					# dump_file = newDir + os.sep + each
					# break
			# if os.path.exists(dump_file):
				# d["dump"] = fileToBuf(dump_file)
				# ret = checkDetectionTricks(dump_file)
				# d["tricks_vm"] = ret
			# else:
				# doLog("No dump file")
				# d["dump"] = None
				# ret = checkDetectionTricks(malwareFile)
				# d["tricks_vm"] = ret
					
			# dump file generated by procdump will cause seriously false positive in checkDetectionTricks()
			d["dump"] = None
			ret = checkDetectionTricks(malwareFile)
			d["tricks_vm"] = ret
				
			doLog("Finish dumping")
			return d
		except Exception as e:
			exc_type, exc_value, exc_traceback =  sys.exc_info()
			traceback.print_exception(exc_type, exc_value, exc_traceback)
			os.system(PM_PATH + " /terminate")
			raise e
		finally:
			self.busy = False

	def ping(self):
		if self.busy:
			return "[BUSY]"
		else:
			return "[ALIVE]"

def usage():
	print "Usage:", sys.argv[0], "[port]"
	print
			
def main(port = 8000):
	d = CXmlRpcDumper()
	server = AsyncXMLRPCServer(('localhost', port), allow_none = True)
	doLog("Listening on port %d..." % port)
	server.register_instance(d)
	try:
		print 'Use Control-C to exit'
		server.serve_forever()
	except KeyboardInterrupt:
		print 'Exiting'

if __name__ == "__main__":
	if len(sys.argv) > 2:
		usage()
	elif len(sys.argv) == 2:
		port = int(sys.argv[1])
		main(port)
	else:
		main()
